import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from python_ai.common.xcommon import *

n_steps = 20
x = tf.range(1, n_steps * 4 + 1)
x = tf.reshape(x, [-1, n_steps])
print(tf.shape(x).numpy())
tf.print(x)

xmax = tf.reduce_max(x)
l1 = layers.Embedding(input_dim=xmax + 1, output_dim=64)
l2 = layers.LSTM(128)
l3 = layers.Dense(10)

sep('1: embedding')  # [ 4 20] => [ 4 20 64]
x = l1(x)
print(tf.shape(x).numpy())
tf.print(x)

sep('2: LSTM')  # [ 4 20 64] => [  4 128]
x = l2(x)
print(tf.shape(x).numpy())
tf.print(x)

sep('3: Dense')  # [  4 128] => [ 4 10]
x = l3(x)
print(tf.shape(x).numpy())
tf.print(x)
